package com.atguigu.study.util;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @auther zzyy
 * @create 2024-12-04 17:40
 */
@Slf4j
public class PerformanceRunnerUtil
{
    //1 静态内部类，封装返回压测结果集PerformanceRunnerResult
    @Data
    public static class PerformanceRunnerResult
    {
        // 请求总数
        private int requests;
        // 并发量
        private int concurrency;
        // 成功请求数
        private int successRequests;
        // 失败请求数
        private int failRequests;
        // 请求总耗时(ms)
        private int timeTakenForTests;
        // 每秒请求数（吞吐量）
        private float requestsPerSecond;
        // 每个请求平均耗时(ms)
        private float timePerRequest;
        // 最快的请求耗时(ms)
        private float fastestCostTime;
        // 最慢的请求耗时(ms)
        private float slowestCostTime;
    }


    //2 定义一个通用的压测方法
    /**
     * @param requests      总请求数
     * @param concurrency   并发数量
     * @param runnable      需要执行的压测业务逻辑代码，直接作为线程任务开干
     * @return              压测结果 {@link PerformanceRunnerResult}
     * @param <T>
     * @throws InterruptedException
     */
    public static <T> PerformanceRunnerResult performanceRun(int requests, int concurrency, Runnable runnable) throws InterruptedException
    {
        log.info("PerformanceRunnerUtil压测开始......");
        PerformanceRunnerResult retValue = new PerformanceRunnerResult();

        //2.1 创建线程池并将所有核心线程池都准备好core=max，直接干到最大
        ThreadPoolExecutor threadPool = new ThreadPoolExecutor(
                concurrency,
                concurrency,
                0L,
                TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>());
        //2.2 启动所有核心线程
        threadPool.prestartAllCoreThreads();
        //2.3 创建一个 CountDownLatch，用于阻塞当前线程池待所有请求处理完毕后，让当前线程继续向下走
        CountDownLatch countDownLatch = new CountDownLatch(requests);
        //2.4 成功请求数-最快耗时-最慢耗时
        AtomicInteger successRequests = new AtomicInteger();
        AtomicInteger fastestCostTime = new AtomicInteger(Integer.MAX_VALUE);
        AtomicInteger slowestCostTime = new AtomicInteger(Integer.MIN_VALUE);


        long startTime = System.currentTimeMillis();
        try
        {
            for (int i = 0; i < requests; i++)
            {
                threadPool.execute(() -> {
                    try
                    {
                        long requestStartTime = System.currentTimeMillis();
                        //执行被压测的方法
                        runnable.run();
                        //业务逻辑压测方法执行耗时
                        int costTime = (int) (System.currentTimeMillis() - requestStartTime);
                        //请求最快耗时
                        setFastestCostTime(fastestCostTime, costTime);
                        //请求最慢耗时
                        setSlowestCostTimeCostTime(slowestCostTime, costTime);
                        //成功请求数+1
                        successRequests.incrementAndGet();
                    } catch (Exception e) {
                        log.error(e.getMessage());
                    } finally {
                        countDownLatch.countDown();
                    }
                });
            }
            //阻塞当前线程，等到压测结束后，该方法会被唤醒，线程继续向下走
            countDownLatch.await();
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            //关闭线程池
            threadPool.shutdown();
        }
        long endTime = System.currentTimeMillis();
        System.out.println("----costTime: "+(endTime - startTime) +" 毫秒");

        //组装最后的结果并返回
        retValue.setRequests(requests);
        retValue.setConcurrency(concurrency);
        retValue.setSuccessRequests(successRequests.get());
        retValue.setFailRequests(requests - retValue.getSuccessRequests());
        retValue.setTimeTakenForTests((int) (endTime - startTime));
        retValue.setRequestsPerSecond((float) requests * 1000f / (float) (retValue.getTimeTakenForTests()));
        retValue.setTimePerRequest((float) retValue.getTimeTakenForTests() / (float) requests);
        retValue.setFastestCostTime(fastestCostTime.get());
        retValue.setSlowestCostTime(slowestCostTime.get());

        return retValue;
    }

    private static void setFastestCostTime(AtomicInteger fastestCostTime, int costTime) {
        while (true) {
            int fsCostTime = fastestCostTime.get();
            if (fsCostTime < costTime) {
                break;
            }
            if (fastestCostTime.compareAndSet(fsCostTime, costTime)) {
                break;
            }
        }
    }

    private static void setSlowestCostTimeCostTime(AtomicInteger slowestCostTime, int costTime) {
        while (true) {
            int slCostTime = slowestCostTime.get();
            if (slCostTime > costTime) {
                break;
            }
            if (slowestCostTime.compareAndSet(slCostTime, costTime)) {
                break;
            }
        }
    }
}